# 数据加载
import os
import numpy as np
import torch
from matplotlib import pyplot as plt
from torchvision import datasets, transforms
from torch.utils.data import DataLoader as TDataLoader, TensorDataset
from torchvision.utils import make_grid
import pandas_datareader as pdr


class DataLoader:
    pass

    def get_loader(self, train_dataset, test_dataset, batch_size):
        train_loader = TDataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        test_loader = TDataLoader(test_dataset, batch_size=1000, shuffle=False)
        return train_loader, test_loader


def data_preprocessing(mean, std):
    # 数据预处理：转换为Tensor并归一化
    transform = transforms.Compose([
        # 将PIL图像或numpy数组转换为PyTorch张量
        transforms.ToTensor(),
        # 归一化到[-1,1]范围
        transforms.Normalize(mean, std)
    ])
    return transform


def get_mnist_loader(batch_size=64):
    # MNIST的均值和标准差
    transform = data_preprocessing((0.1307,), (0.3081,))
    # 下载并加载训练集和测试集
    train_dataset = datasets.MNIST(root=ROOT_PATH, train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST(root=ROOT_PATH, train=False, download=True, transform=transform)

    # 创建DataLoader（批量加载）
    train_loader = TDataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = TDataLoader(test_dataset, batch_size=1000, shuffle=False)
    return train_loader, test_loader


def get_cifar_10_loader(batch_size=64):
    transform = data_preprocessing((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    train_dataset = datasets.CIFAR10(root=ROOT_PATH, train=True, download=True, transform=transform)
    test_dataset = datasets.CIFAR10(root=ROOT_PATH, train=False, download=True, transform=transform)
    train_loader = TDataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = TDataLoader(test_dataset, batch_size=1000, shuffle=False)
    return train_loader, test_loader


def get_flowers_102_loader(batch_size=256):
    transform_train = transforms.Compose([
        transforms.RandomRotation(30),
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomVerticalFlip(p=0.5),
        # 将PIL图像或numpy数组转换为PyTorch张量
        transforms.ToTensor(),
        # 归一化到[-1,1]范围
        transforms.Normalize((0.458, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])

    transform_test = transforms.Compose([
        transforms.RandomResizedCrop(224),
        # 将PIL图像或numpy数组转换为PyTorch张量
        transforms.ToTensor(),
        # 归一化到[-1,1]范围
        transforms.Normalize((0.458, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])

    train_dataset = datasets.Flowers102(root=ROOT_PATH, split="test", download=True, transform=transform_train)
    test_dataset = datasets.Flowers102(root=ROOT_PATH, split="train", download=True, transform=transform_test)
    train_loader = TDataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = TDataLoader(test_dataset, batch_size=256, shuffle=False)
    return train_loader, test_loader


def get_gs10_loader(batch_size=256):
    gs10 = pdr.get_data_fred("GS10")
    # 数量
    num = len(gs10)
    # 价格列表
    x = torch.tensor(gs10["GS10"].to_list())
    # 预测序列长度
    seq_len = 8
    # #设置批大小
    batch_size = 2
    # 全零初始化特征矩阵，mum - seq_len行，seq_len列
    x_feature = torch.zeros((num - seq_len, seq_len))
    y_label = torch.zeros((num - seq_len, seq_len))
    for i in range(seq_len):
        x_feature[:, i] = x[i:num - seq_len + i]  # 为特征矩阵赋值
        y_label[:, i] = x[i + 1:num - seq_len + i + 1]  # 真实结果列表

    x_ = x_feature[: num - seq_len].unsqueeze(2)
    y_ = y_label[:num - seq_len]
    train_loader = TDataLoader(
        TensorDataset(x_, y_),
        batch_size=batch_size,
        shuffle=True)  # 构建数据加载器

    return train_loader


def get_loader(name, num_workers=1, batch_size=64):
    if name == "mnist":
        train_loader, test_loader = get_mnist_loader(batch_size=batch_size)
    elif name == "cifar10":
        train_loader, test_loader = get_cifar_10_loader(batch_size=batch_size)
    elif name == "flowers102":
        train_loader, test_loader = get_flowers_102_loader(batch_size=batch_size)
    elif name == "gs10":
        train_loader = get_gs10_loader(batch_size=batch_size)
        test_loader = train_loader
    return train_loader, test_loader


def imshow(img):
    img = img / 2 + 0.5  # 反归一化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


def show_data(name):
    # 4. 定义类别名称（与CIFAR-10一致）
    classes = ('plane', 'car', 'bird', 'cat', 'deer',
               'dog', 'frog', 'horse', 'ship', 'truck')

    if name == 'mnist':
        train_loader, test_loader = get_mnist_loader()
    elif name == 'cifar10':
        train_loader, test_loader = get_cifar_10_loader()
    # 获取一个批次的训练图像
    dataiter = iter(train_loader)
    images, labels = next(dataiter)

    # 显示图像
    imshow(make_grid(images))
    # 打印标签
    print(' '.join(f'{classes[labels[j]]:5s}' for j in range(64)))


if __name__ == '__main__':
    # show_data("cifar10")
    get_gs10_loader()
